/**
 * @file   MultigridSolver.h
 * @author HirasawaYui <yui@Ubuntu18-04>
 * @date   Wed Nov 18 14:32:50 2020
 * 
 * @brief  
 * 
 * 
 */

#include <iostream>
#include <Eigen/Sparse>
#include <vector>
#include "Mesh.h"

#define Multi2(a,bit) (a<<bit)
#define Devide2(a,bit) (a>>bit)
#define POW2(a) (1<<a)
#define TEMPLATE template<Dimension DIM,typename MatrixType>
typedef double Real;
typedef Eigen::VectorXd VectorXd;
TEMPLATE
class MultigridSolver
{
public:
    typedef MatrixType SpMat;
    typedef Eigen::Triplet<Real> Tri;
    typedef double(*Function)(double*); 
protected:
    SpMat _A_origin;
    SpMat _A_temp;
    VectorXd _Rhs_origin;
    VectorXd _Rhs_temp;
    std::vector<VectorXd> _Rhs;
    std::vector<VectorXd> _u;
    Function _BCF;
    Function _F;
    int _nowlevel;
    Mesh<DIM>* _FineMesh;
    Mesh<DIM>* _CoarseMesh;
    Eigen::ConjugateGradient<MatrixType > Solver_sparse;
    Real _Tolerance = 1e-6;
    long _MaxIternum =  20; 
    /// the transfer matrix
    SpMat _P;
public:
    virtual VectorXd& Solve() = 0;
    virtual void setTolerance(Real tolerance) = 0;
};

#undef TEMPLATE
#define TEMPLATE template<typename MatrixType>
TEMPLATE
class Q1_MGSolver:public MultigridSolver<2,MatrixType>
{
public:
    typedef MultigridSolver<2,MatrixType> _base;
    typedef typename _base::SpMat SpMat;
    typedef typename _base::Function Function;
    typedef Eigen::Triplet<Real> Tri;
    using _base::_A_origin;
    using _base::_A_temp;
    using _base::_Rhs_origin;
    using _base::_Rhs_temp;
    using _base::_Rhs;
    using _base::_u;
    using _base::_BCF;
    using _base::_F;
    using _base::_nowlevel;
    using _base::_FineMesh;
    using _base::_CoarseMesh;
    using _base::Solver_sparse;
    using _base::_P;
    using _base::_Tolerance;
    using _base::_MaxIternum;
public:
    Q1_MGSolver(SpMat A, VectorXd Rhs, Function F, Function BCF, Mesh<2>* mesh);
    /** 
     * @brief Generate transfer matrix P
     * 
     */
    ~Q1_MGSolver()
    {
        delete _FineMesh;
        delete _CoarseMesh;
    }
    void GenerateP();
    SpMat P() const {return _P;};
    void DealWithBoundaryCondition(SpMat &A, VectorXd &Rhs);
    void UpLevel();
    void DownLevel();
    void VCycle();
    void Smooth(SpMat A, VectorXd Rhs);
    VectorXd& Solve();
    void Reset();
    void compute(SpMat A){_A_origin= A;_A_temp = A;};
    void setTolerance(Real tolerance){_Tolerance = tolerance;};
};
TEMPLATE
Q1_MGSolver<MatrixType>::Q1_MGSolver(SpMat A, VectorXd Rhs, Function F, Function BCF, Mesh<2>* mesh)
{
    _A_origin = A;
    _A_temp = A;
    _Rhs_origin = Rhs;
    _Rhs_temp = Rhs;
    _Rhs.push_back(Rhs);
    _u.push_back(Eigen::MatrixXd::Zero(Rhs.size(),1));
    _F = F;
    _BCF = BCF;
    _nowlevel = 1;
    _FineMesh = new Q1Mesh(mesh->RDomain(), {mesh->x_seg(), mesh->y_seg()});
    _CoarseMesh = new Q1Mesh(_FineMesh->RDomain(), {Devide2(_FineMesh->x_seg(),1), Devide2(_FineMesh->y_seg(),1)});
    Solver_sparse.setTolerance(1e-6);
}
TEMPLATE
void Q1_MGSolver<MatrixType>::GenerateP()
{
    _P = SpMat(_FineMesh->n_point(), _CoarseMesh->n_point());
    std::vector<Tri> TriList(_FineMesh->n_point()* _CoarseMesh->n_point());
    std::vector<Tri>::iterator it = TriList.begin();
    for (int i = 0; i < _CoarseMesh->y_seg()+1; i++)
	for (int j = 0; j < _CoarseMesh->x_seg()+1; j++)
	{
	    int CoarseIdx = _CoarseMesh->CortoIdx(i, j);
	    int FineIdx = _FineMesh->CortoIdx(2*i, 2*j);
        //if(_CoarseMesh->IsBoundary(CoarseIdx))
	    *(it++) = Tri(FineIdx, CoarseIdx, 1.0);
        //else
        //{
        if (i != 0)
	    {
		*(it++) = Tri(FineIdx-(_FineMesh->x_seg()+1), CoarseIdx, 0.5);
		if (j != 0)
		    *(it++) = Tri(FineIdx-(_FineMesh->x_seg()+1)-1, CoarseIdx, 0.25);
		if (j != _CoarseMesh->x_seg())
		    *(it++) = Tri(FineIdx-(_FineMesh->x_seg()+1)+1, CoarseIdx, 0.25);
	    }
	    if (i != _CoarseMesh->y_seg())
	    {
		*(it++) = Tri(FineIdx+(_FineMesh->x_seg()+1), CoarseIdx, 0.5);
		if (j != 0)
		    *(it++) = Tri(FineIdx+(_FineMesh->x_seg()+1)-1, CoarseIdx, 0.25);
		if (j != _CoarseMesh->x_seg())
		    *(it++) = Tri(FineIdx+(_FineMesh->x_seg()+1)+1, CoarseIdx, 0.25);
	    }
	    if (j != 0)
		*(it++) = Tri(FineIdx-1, CoarseIdx, 0.5);
	    if (j != _CoarseMesh->x_seg())
		*(it++) = Tri(FineIdx+1, CoarseIdx, 0.5);
        //}
	}
    _P.setFromTriplets(TriList.begin(), TriList.end());
    _P.makeCompressed();
}

TEMPLATE
void Q1_MGSolver<MatrixType>::DealWithBoundaryCondition(SpMat &A, VectorXd &Rhs)
{
    for (auto k:_FineMesh->Boundary())
    {
	Dofs<2> bnd_point = _FineMesh->DofsofIndex(k);
	Real bnd_value = _BCF(*bnd_point);
    if(_nowlevel == 1)
	    Rhs[k] = bnd_value * A.coeffRef(k, k);
	for(Eigen::SparseMatrix<Real>::InnerIterator it(A,k);it;++it)
        {
            int row = it.row();
            if(row == k)
                continue;
            if(_nowlevel == 1)
                Rhs[row] -= A.coeffRef(k,row)* bnd_value;
            A.coeffRef(k,row) = 0.0;
            A.coeffRef(row,k) = 0.0;
        }
    }
}
TEMPLATE
void Q1_MGSolver<MatrixType>::UpLevel()
{
    _nowlevel++;
    _FineMesh = _CoarseMesh;
    _CoarseMesh = new Q1Mesh(_FineMesh->RDomain(), {Devide2(_FineMesh->x_seg(),1), Devide2(_FineMesh->y_seg(),1)});
    _A_temp = _P.transpose() * _A_temp * _P;
    //std::cout << _A_temp << std::endl;
    _Rhs_temp = _Rhs[_nowlevel-1];
}
TEMPLATE
void Q1_MGSolver<MatrixType>::DownLevel()
{
    _nowlevel--;
    _CoarseMesh = _FineMesh;
    _FineMesh = new Q1Mesh(_CoarseMesh->RDomain(), {Multi2(_CoarseMesh->x_seg(),1), Multi2(_CoarseMesh->y_seg(),1)});
}
TEMPLATE
void Q1_MGSolver<MatrixType>::Smooth(SpMat A, VectorXd Rhs)
{
    VectorXd& u = _u[_nowlevel-1];
    VectorXd u_star = Eigen::MatrixXd::Zero(Rhs.size(),1);
    for(int i = 0;i < u.size();i++)
    {
        if(_FineMesh->IsBoundary(i))
        {
            Dofs<2> bnd_point = _FineMesh->DofsofIndex(i);
            u_star[i] = _BCF(*bnd_point);
        }
        else
        {
             double a = Rhs[i];
            for(Eigen::SparseMatrix<Real>::InnerIterator it(A,i);it;++it)
            {
                if(it.row() == i)
                    continue;
                else
                {
                    a -=(it.value() * (u[it.row()]));
                }
            }
            u_star[i] = a / A.coeffRef(i,i);
        }  
        
    }
    for(int i = 0; i < u.size();i++)
        u[i] = (1.0 - 1.0)* u[i] + 1.0 * u_star[i];
    _u[_nowlevel-1] = u;
    //std::cout << " u is :" << u <<std::endl;
}
TEMPLATE
void Q1_MGSolver<MatrixType>::VCycle()
{
    SpMat A = _A_temp;
    SpMat A_dealt = A;
    VectorXd Rhs = _Rhs_temp;
    VectorXd Rhs_dealt = _Rhs[_nowlevel - 1];
    //std::cout <<"边界条件处理前A_dealt："<<A_dealt << std::endl;
    //std::cout <<"边界条件处理前的Rhs_dealt"<<Rhs_dealt << std::endl;
    DealWithBoundaryCondition(A_dealt, Rhs_dealt);
    //std::cout <<"边界条件处理之后A_dealt："<<A_dealt << std::endl;
    //std::cout <<"边界条件处理之后的Rhs_dealt"<<Rhs_dealt << std::endl;
    GenerateP();
    if (_nowlevel == 3)
	{
        Solver_sparse.compute(A_dealt);
        _u[_nowlevel - 1] = Solver_sparse.solve(Rhs_dealt);
    } 
    else
    {
	for (int i = 0; i < 5; i++)
	    Smooth(A_dealt,Rhs_dealt);
    //std::cout << " f- av" << (Rhs_dealt - A_dealt * _u[_nowlevel-1])<<std::endl;
	VectorXd r_2h = _P.transpose()*(Rhs_dealt - A_dealt * _u[_nowlevel-1]);
    for(int i = 0;i < r_2h.size();i++)
        if(_CoarseMesh->IsBoundary(i))
            r_2h[i] = 0.0;
    //std::cout << "r_2h" << r_2h << std::endl;
 	VectorXd e = Eigen::MatrixXd::Zero(r_2h.size(),1);
	_u.push_back(e);
	_Rhs.push_back(r_2h);
	UpLevel();
	VCycle();
	DownLevel();
	GenerateP();
	_u[_nowlevel-1] = _u[_nowlevel-1] + _P * _u[_nowlevel];
    //std::cout << _u[_nowlevel - 1] << std::endl;
	for (int i = 0; i < 5; i++)
	    Smooth(A_dealt,Rhs_dealt);
    }
    //std::cout << _u[0] <<std::endl;
}
TEMPLATE
void Q1_MGSolver<MatrixType>::Reset()
{
    _A_temp = _A_origin;
    _Rhs_temp = _Rhs_origin;
    for(int i = 1; i < _Rhs.size();i++)
    {
        _Rhs[i].setZero();
        _u[i].setZero();
    }
}
TEMPLATE
VectorXd& Q1_MGSolver<MatrixType>::Solve()
{
    int times = 0;
    VectorXd residual_vector = _A_origin * Eigen::MatrixXd::Zero(_Rhs[0].size(),1) - _Rhs_origin;
    double residual = residual_vector.squaredNorm();
    double residual_max = residual_vector.lpNorm<Eigen::Infinity>();
    do
    {

        VCycle();
        Reset();
        residual_vector = _A_origin * _u[0] - _Rhs_origin;
        //std::cout << " residual _vector" << residual_vector << std::endl;
        residual = residual_vector.squaredNorm();
        residual_max = residual_vector.lpNorm<Eigen::Infinity>();
        std::cout << " the "<< times + 1 << "th l2 norm of residual vecror is :" << residual << " and the max norm case is :" << residual_max<<std::endl;
    }while(times++ < 100& residual >= _Tolerance);
    return _u[0];
}
#undef TEMPLATE

